import numpy as np
import cv2 as cv


METHOD_DIFF_CHANNEL = 0
METHOD_HSV_RANGE    = 1
METHOD_HSV_RANGE2   = 2
METHOD_HSV_RANGE3   = 3

def searchRedLaserPoint(img,method=0):
    # 输入BGR图像，所以要重新写一下
    if method == METHOD_DIFF_CHANNEL:
        diffRGImage = np.uint8(np.array(img[:,:,2],dtype=np.int16)-np.array(img[:,:,1],dtype=np.int16) > 75)
        # diffGRImage = np.uint8(np.array(img[:,:,1],dtype=np.int16)-np.array(img[:,:,0],dtype=np.int16) > 45)
        ksize = 11

        while ksize > 1:
            filteredImage = cv.morphologyEx(diffRGImage,cv.MORPH_OPEN,np.ones((ksize,ksize)))
            imgMomentsL = cv.moments(filteredImage,False)
            # 接下来就是求矩
            if imgMomentsL['m00'] == 0:
                ksize -= 2
                continue
            else:
                return imgMomentsL['m10']/imgMomentsL['m00'],imgMomentsL['m01']/imgMomentsL['m00']
        return None
    elif method == METHOD_HSV_RANGE:
        redLaserHSVLower1 = np.array([120,0,200])
        redLaserHSVUpper1 = np.array([180,255,255])
        redLaserHSVLower2 = np.array([120,0,20])
        redLaserHSVUpper2 = np.array([180,255,100])
        imgHSV = cv.cvtColor(img,cv.COLOR_BGR2HSV)
        # 直方图均衡化
        # imgHSV[:,:,1] = cv.equalizeHist(imgHSV[:,:,1])
        imgHSV[:,:,2] = cv.equalizeHist(imgHSV[:,:,2])

        imgMask1 = cv.inRange(imgHSV,redLaserHSVLower1,redLaserHSVUpper1)
        imgMask2 = cv.inRange(imgHSV,redLaserHSVLower2,redLaserHSVUpper2)
        imgMask = cv.bitwise_or(imgMask1,imgMask2)
        kernel1 = np.ones((3,3))
        imgMask = cv.morphologyEx(imgMask,cv.MORPH_CLOSE,kernel1)
        imgMask = cv.morphologyEx(imgMask,cv.MORPH_OPEN,kernel1)
        imgMomentsL = cv.moments(imgMask,False)
        if imgMomentsL['m00'] == 0:
            return None
        else:
            return imgMomentsL['m10']/imgMomentsL['m00'],imgMomentsL['m01']/imgMomentsL['m00']
    elif method == METHOD_HSV_RANGE2 or method == METHOD_HSV_RANGE3:
        if method == METHOD_HSV_RANGE2:
            redLaserHSVLower1 = np.array([165,50,150])
            redLaserHSVUpper1 = np.array([179,255,255])
            redLaserHSVLower2 = np.array([165,50,150])
            redLaserHSVUpper2 = np.array([179,255,255])
        else:
            redLaserHSVLower1 = np.array([165,115,70])
            redLaserHSVUpper1 = np.array([179,255,255])
            redLaserHSVLower2 = np.array([165,115,70])
            redLaserHSVUpper2 = np.array([179,255,255])
        imgHSV = cv.cvtColor(img,cv.COLOR_BGR2HSV)
        imgMask1 = cv.inRange(imgHSV,redLaserHSVLower1,redLaserHSVUpper1)
        imgMask2 = cv.inRange(imgHSV,redLaserHSVLower2,redLaserHSVUpper2)
        imgMask = cv.bitwise_or(imgMask1,imgMask2)
        kernel = np.ones((5,5))

        imgMask = cv.morphologyEx(imgMask,cv.MORPH_CLOSE,kernel)
        imgMask = cv.morphologyEx(imgMask,cv.MORPH_OPEN,kernel)
        
        contours,hierarchy = cv.findContours(imgMask,cv.RETR_CCOMP,cv.CHAIN_APPROX_SIMPLE)
        # print(hierarchy)
        # 寻找最大的轮廓
        maxAreaContour = None
        maxAreaContourIndex = 0
        maxArea = 0
        for i in range(len(contours)):
            c = contours[i]
            area = cv.contourArea(c)
            if area > maxArea:
                maxAreaContourIndex = i
                maxAreaContour = c
                maxArea = area
        if type(maxAreaContour) == type(None):
            return None
        else:
            # 寻找firstChild
            fcIndex = hierarchy[0][maxAreaContourIndex][2]
            if fcIndex != -1:
                return cv.minEnclosingCircle(contours[fcIndex])[0]
            else:
                return cv.minEnclosingCircle(contours[maxAreaContourIndex])[0]
        

def searchGreenLaserPoint(img,method=0):
    greenLaserHSVLower1 = np.array([70,40,220])
    greenLaserHSVUpper1 = np.array([105,255,255])
    greenLaserHSVLower2 = np.array([50,0,225])
    greenLaserHSVUpper2 = np.array([115,255,255])
    imgHSV = cv.cvtColor(img,cv.COLOR_BGR2HSV)
    # 直方图均衡化
    # imgHSV[:,:,1] = cv.equalizeHist(imgHSV[:,:,1])
    # imgHSV[:,:,2] = cv.equalizeHist(imgHSV[:,:,2])

    imgMask1 = cv.inRange(imgHSV,greenLaserHSVLower1,greenLaserHSVUpper1)
    imgMask2 = cv.inRange(imgHSV,greenLaserHSVLower2,greenLaserHSVUpper2)
    imgMask = cv.bitwise_or(imgMask1,imgMask2)
    kernel = np.ones((3,3))
    imgMask = cv.morphologyEx(imgMask,cv.MORPH_CLOSE,kernel)
    imgMask = cv.morphologyEx(imgMask,cv.MORPH_OPEN,kernel)
    imgMomentsL = cv.moments(imgMask,False)

    if imgMomentsL['m00'] == 0:
        return None
    else:
        return imgMomentsL['m10']/imgMomentsL['m00'],imgMomentsL['m01']/imgMomentsL['m00']